import numpy as np
from matplotlib import pyplot as plt

from src.core.NeuralNetwork import NeuralNetwork


def simulate_snn():
    # 创建网络
    snn = NeuralNetwork(n_input=2, n_hidden=4, n_output=1)

    # 模拟参数
    time_steps = 100
    dt = 1.0
    input_pattern = np.zeros((time_steps, 2))

    # 生成简单的输入模式（前50步一个神经元活跃，后50步另一个活跃）
    input_pattern[:50, 0] = 1.0
    input_pattern[50:, 1] = 1.0

    # 记录输出
    output_history = []

    # 运行仿真
    for t in range(time_steps):
        inputs = input_pattern[t]
        outputs = snn.forward(inputs, dt)
        output_history.append(outputs[0])

    # 可视化结果
    plt.figure(figsize=(12, 6))

    plt.subplot(2, 1, 1)
    plt.title("Input Spikes")
    plt.imshow(input_pattern.T, aspect='auto', cmap='binary')
    plt.xlabel("Time")
    plt.ylabel("Input Neuron")

    plt.subplot(2, 1, 2)
    plt.title("Output Spikes")
    plt.plot(output_history, 'b-', label='Output Neuron')
    plt.xlabel("Time")
    plt.ylabel("Spike")
    plt.legend()

    plt.tight_layout()
    plt.show()


# 运行模拟
if __name__ == "__main__":
    simulate_snn()
